import torch
from torch import nn
from FFN import FFN


class Expert(nn.Module):

    def __init__(self, num_experts, top_k, input_dim, hide_dim):
        super().__init__()

        self._top_k = top_k
        self._experts = nn.ModuleList(
            [FFN(input_dim, hide_dim) for _ in range(num_experts)]
        )
        self._gate = nn.Linear(input_dim, num_experts)

    def forward(self, x):
        print(x.shape)
        _bn, _seq, _vec = x.shape
        print(_bn, _seq, _vec)
        _gate_out = self._gate(x)
        print(_gate_out.shape)
        _top_values, _top_indices = torch.topk(_gate_out, self._top_k, dim=-1)
        print(_top_values, _top_indices)
        _output = torch.zeros(_bn, _seq, self._top_k, _vec).to(x.device)
        print(_output.shape)

        for _i in range(_bn):
            for _j in range(_seq):
                for _k in range(self._top_k):
                    _expert = self._experts[_top_indices[_i, _j, _k]]
                    _output[_i, _j, _k] = _expert(x[_i, _j])

        # return _output.mean(dim=2)

        w = torch.softmax(_top_values, dim=-1)[:, :, :, None]
        return (w * _output).sum(dim=2)
